#include "yolov8_pose.h"
// #include <opencv2/opencv.hpp>

Yolov8Pose::Yolov8Pose()
{
}

Yolov8Pose::~Yolov8Pose()
{
}

std::vector<PoseInfo> Yolov8Pose::Inference(cv::Mat input_mat, float scale,
                                            std::vector<int> &pad,
                                            float score_threshold,
                                            float iou_threshold,
                                            float nms_threshold)

{
    std::vector<PoseInfo> result;

    // -----------预处理----------
    int channels = input_mat.channels();
    int height = input_mat.rows;
    int width = input_mat.cols;

    std::vector<float> input_data;
    input_data.resize(channels * height * width);

    for (int c = 0; c < channels; c++)
    {
        for (int h = 0; h < height; h++)
        {
            for (int w = 0; w < width; w++)
            {
                float pix = input_mat.ptr<uchar>(h)[w * 3 + c];
                int chw_idex = c * height * width + h * width + w;
                float tmp = pix / 255.0f;
                input_data[chw_idex] = tmp;
            }
        }
    }

    // ---------创建输入tensor-------
    Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
    std::vector<int64_t> input_node_dims = {1, channels, height, width};
    Ort::Value input_tensor = Ort::Value::CreateTensor<float>(
        memory_info,
        input_data.data(),
        input_data.size(),
        input_node_dims.data(),
        input_node_dims.size());

    assert(input_tensor.IsTensor());

    // -------获取输出temsor--------

    std::vector<const char *> inputNodeNames = {"images"};
    std::vector<const char *> outputNodeNames = {"output0"};

    auto output_tensors = m_session->Run(
        Ort::RunOptions{nullptr},
        inputNodeNames.data(),
        &input_tensor,
        1,
        outputNodeNames.data(),
        outputNodeNames.size());

    // --------后处理-------------
    std::vector<cv::Rect> boxes;
    std::vector<float> scores;
    std::vector<int> class_id;
    std::vector<std::vector<PosePoint>> object_keypoints;

    auto output_data = output_tensors[0].GetTensorMutableData<float>();
    auto output_dims = output_tensors[0].GetTensorTypeAndShapeInfo().GetShape(); // 获取维度

    int64_t rows = output_dims[1]; // 56
    int64_t clos = output_dims[2]; // 8400

    cv::Mat rawData = cv::Mat(rows, clos, CV_32F, output_data).t();

    // output_dims[1]: [x, y, w, h] + score + [17, 3] keypoints
    for (int i = 0; i < clos; i++)
    {
        float class_score = rawData.at<float>(i, 4);

        if (class_score > score_threshold)
        {
            // std::cout << class_score << std::endl;
            scores.push_back(class_score);
            class_id.push_back(0);

            //----------[x, y, w, h]------------
            float x = rawData.at<float>(i, 0);
            float y = rawData.at<float>(i, 1);
            float w = rawData.at<float>(i, 2);
            float h = rawData.at<float>(i, 3);

            //----------映射到原图坐标-----------
            int left = int((x - pad[0] - 0.5 * w) * scale);
            int top = int((y - pad[1] - 0.5 * h) * scale);
            int width = int(w * scale);
            int height = int(h * scale);

            boxes.push_back(cv::Rect(left, top, width, height));

            //---------获取关键点坐标-----------

            std::vector<PosePoint> keypoints;
            cv::Mat kpts = rawData.row(i).colRange(5, 56);

            for (int i = 0; i < 17; i++)
            {
                PosePoint pose;
                float x = (kpts.at<float>(0, i * 3 + 0) - pad[0]) * scale;
                float y = (kpts.at<float>(0, i * 3 + 1) - pad[1]) * scale;
                float v = kpts.at<float>(0, i * 3 + 2);

                pose.x = x;
                pose.y = y;
                pose.visibility = v;

                keypoints.push_back(pose);
            }

            object_keypoints.push_back(keypoints);
        }
    }

    // ------------NMS-------------
    std::vector<int> indices;
    cv::dnn::NMSBoxes(boxes, scores, score_threshold, nms_threshold, indices);
    // std::cout << "indices" << indices.size() << std::endl;

    for (int i = 0; i < indices.size(); i++)
    {
        PoseInfo pose_point;
        int idx = indices[i];

        //---------矩形框信息-----------
        pose_point.bbox.x = boxes[idx].x;
        pose_point.bbox.y = boxes[idx].y;
        pose_point.bbox.w = boxes[idx].width;
        pose_point.bbox.h = boxes[idx].height;
        pose_point.bbox.score = scores[idx];
        pose_point.bbox.label = class_id[idx];
        pose_point.bbox.label_name = "person";

        std::vector<PosePoint> kpts = object_keypoints[idx];
        for (const auto &kpt : kpts)
        {
            pose_point.poseV.push_back(kpt);
        }
        result.push_back(pose_point);
    }

    return result;
}
